package com.captjack.math;

import java.math.BigInteger;

/**
 * @author Jack Sparrow
 * @version 1.0.0
 * @date 2022/10/16 21:55
 * package com.captjack.docker.simple.application
 */
public class MatrixQuickPower {

    /**
     * 矩阵乘法
     *
     * @param a   a
     * @param b   b
     * @param mod mod
     * @return multiply
     */
    public static long[][] multiply(long[][] a, long[][] b, long mod) {
        int rowA = a.length, colA = a[0].length, colB = b[0].length;
        long[][] resultData = new long[rowA][colA];
        for (int i = 0; i < rowA; i++) {
            for (int j = 0; j < colB; j++) {
                for (int k = 0; k < colA; k++) {
                    resultData[i][j] = (resultData[i][j] + (a[i][k] * b[k][j]) % mod) % mod;
                }
            }
        }
        return resultData;
    }

    /**
     * 矩阵速降幂
     *
     * @param a     底数
     * @param power 指数
     * @param mod   mod
     * @return mod
     */
    public static long[][] quickMatrixPower(long[][] a, long power, long mod) {
        int rowA = a.length;
        long[][] data = new long[rowA][rowA];
        // 初始化单位矩阵
        for (int i = 0; i < rowA; i++) {
            data[i][i] = 1;
        }
        while (power > 0) {
            if ((power & 1) == 1) {
                data = multiply(data, a, mod);
            }
            a = multiply(a, a, mod);
            // 右移一位
            power >>>= 1;
        }
        return data;
    }

    /**
     * 矩阵速降幂
     *
     * @param a     底数
     * @param power 指数
     * @param mod   mod
     * @return mod
     */
    public static long[][] quickMatrixPower(long[][] a, BigInteger power, long mod) {
        int rowA = a.length;
        long[][] data = new long[rowA][rowA];
        // 初始化单位矩阵
        for (int i = 0; i < rowA; i++) {
            data[i][i] = 1;
        }
        String powerStr = power.toString();
        BigInteger two = new BigInteger("2");
        while (power.compareTo(BigInteger.ZERO) > 0) {
            if (powerStr.charAt(powerStr.length() - 1) % 2 == 1) {
                data = multiply(data, a, mod);
            }
            power = power.divide(two);
            powerStr = power.toString();
            //
            a = multiply(a, a, mod);
        }
        return data;
    }

}
